# coding:utf-8
__author__ = 'kony'

from base.models import Project,Worker
import shutil
import paramiko
import subprocess
import os, time

class FileModel:
    def __init__(self, project, engine, workers):
        self.project = project
        self.engine = engine
        self.workers = workers
        
    # Establish sftp connection
    def sftp_connect(self, username='', password='', host='', port = 22):
        client = None
        sftp = None
        try:
            client = paramiko.Transport((host, port))
        except Exception as error:
            print(error)
            print("connect failed")
        else:
            try:
                client.connect(username=username, password=password)
            except Exception as error:
                print(error)
            else:
                sftp = paramiko.SFTPClient.from_transport(client)
        return client,sftp

    # Check whether the remote directory exists. If it does not exist, create it. (Only level-1 directories can be created.)
    def create_dir(self, sftp, remote):
        try:
            sftp.chdir(remote)
        except Exception as e:
            sftp.mkdir(remote)
            print("Create a directory on remote sftp：{}".format(remote))

    # Upload to remote, "remote" parameter represents the remote path
    def upload(self, sftp, local, remote):
        # Check whether local is a file, obtain the file name, and concatenate it to remote
        if os.path.isfile(local):
            self.create_dir(sftp, remote)    # Check whether the remote directory exists
            filename = os.path.basename(os.path.normpath(local))
            remote = os.path.join(remote, filename)
            try:
                sftp.put(local, remote)
            except Exception as error:
                print(error)
        else:
            print("%s is no file", local)

    def sync(self, records):
        for record in records:
            local = record['path']
            if len(self.workers) == 0:
                msg = {"code": 1, "msg": "The load machine is not configured", "data": ""}
                return msg
            if self.engine == 'jmeter':
                remote = f'/test/tools/apache-jmeter-5.4/data/{self.project}'
                ips = self.workers.split(',')
                for ip in ips:
                    account = Worker.objects.filter(project=self.project, ip=ip).values("username", "password")
                    username = account[0]['username']
                    password = account[0]['password']
                    client, sftp = self.sftp_connect(username=username, password=password, host=ip, port=22)
                    self.upload(sftp, local, remote)
                    client.close()
            elif self.engine == 'locust':
                remote = f'/test/projects/{self.project}'
                ips = self.workers.split(',')
                for ip in ips:
                    account = Worker.objects.filter(project=self.project, ip=ip).values("username", "password")
                    username = account[0]['username']
                    password = account[0]['password']
                    client, sftp = self.sftp_connect(username=username, password=password, host=ip, port=22)
                    self.upload(sftp, local, remote)
                    client.close()
        msg = {"code": 0, "msg": "success", "data": ""}
        return msg

    # Used to split parameter files
    def split(self, records):
        record = records[0]
        local = record['path']                                  # Parameter file to synchronize
        path = os.path.dirname(os.path.normpath(local))         # The relative path to the parameter file
        filename = os.path.basename(os.path.normpath(local))    # File name
        tmpfile = os.path.join(path,'tmp', filename)            # Split temporary cache file
        record['path'] = tmpfile
        records[0] = record
        ips = self.workers.split(',')
        n = 0
        count = subprocess.getoutput("wc -l %s" % local)
        rows = int(int(count.split()[0]) / len(ips))
        with open(local) as fin:
            fout = open(tmpfile, "w")
            for i,line in enumerate(fin, 1):
                fout.write(line)
                if i % rows == 0:
                    fout.close()
                    self.workers = ips[n]
                    n += 1
                    msg = self.sync(records)
                    fout = open(tmpfile, "w")
            fout.close()
        return msg
